package org.liu.knowledge.rag;

import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.injector.ContentInjector;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;

import java.util.List;

public class RagDemo {
    static EmbeddingModel embeddingModel = OllamaEmbeddingModel.builder()
            .baseUrl("http://localhost:11434/")
            .modelName("zyw0605688/gte-large-zh")
            .build();

    public static void main(String[] args) {
        // 向量数据库
        EmbeddingStore<TextSegment> embeddingStore = PgVectorEmbeddingStore.builder()
                //指定主机地址
                .host("192.168.10.100")
                //指定端口
                .port(5432)
                //指定用户名
                .user("root")
                //指定密码
                .password("Password123@postgres")
                //指定数据库名
                .database("postgres")
                //指定向量数据所在表名
                .table("knowledge_vector")
                //指定向量维度
                .dimension(embeddingModel.dimension())
                .createTable(true)
                .build();

        // 向量数据检索器
        EmbeddingStoreContentRetriever retriever = EmbeddingStoreContentRetriever.builder()
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                .maxResults(5)
                .minScore(0.5)
                .build();


        String question = "what is i like ?";
        Query query = Query.from(question);
        List<Content> contents = retriever.retrieve(query);

        contents.forEach(System.out::println);


        //提示模板
        PromptTemplate promptTemplate = new PromptTemplate(
                "用户问题：{{userMessage}}" +
                "使用以下信息回答:" +
                "{{contents}}");
        DefaultContentInjector injector = DefaultContentInjector.builder()
//                .metadataKeysToInclude()
                .promptTemplate(promptTemplate)
                .build();
        // 生成提示词
        UserMessage inject = injector.inject(contents, UserMessage.userMessage(question));
        System.out.println(inject);
    }
}
